/*
 * AIEngine a new generation network intrusion detection system.
 *
 * Copyright (C) 2013-2023  Luis Campo Giralte
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 *
 * Written by Luis Campo Giralte <luis.camp0.2009@gmail.com>
 *
 */
#include "SMBProtocol.h"
#include <iomanip>

namespace aiengine {

SMBProtocol::~SMBProtocol() {

	anomaly_.reset();
}

bool SMBProtocol::is_minimal_smb_header(const uint8_t *hdr) {

	if ((hdr[0] == 0x00)and(hdr[1] == 0x00)) {
		if ((hdr[5] == 'S')and(hdr[6] == 'M')and(hdr[7] == 'B')) {
			return true;
		}
	}
	return false;
}

bool SMBProtocol::check(const Packet &packet) {

	int length = packet.getLength();

	if (length >= 8) {
		setHeader(packet.getPayload());

		if (is_minimal_smb_header(header_)) {
			++total_valid_packets_;
			return true;
		}
	}
	++total_invalid_packets_;
	return false;
}

void SMBProtocol::setDynamicAllocatedMemory(bool value) {

        info_cache_->setDynamicAllocatedMemory(value);
        filename_cache_->setDynamicAllocatedMemory(value);
}

bool SMBProtocol::isDynamicAllocatedMemory() const {

        return info_cache_->isDynamicAllocatedMemory();
}

uint64_t SMBProtocol::getCurrentUseMemory() const {

        uint64_t mem = sizeof(SMBProtocol);

        mem += info_cache_->getCurrentUseMemory();
        mem += filename_cache_->getCurrentUseMemory();

        return mem;
}

uint64_t SMBProtocol::getAllocatedMemory() const {

        uint64_t mem = sizeof(SMBProtocol);

        mem += info_cache_->getAllocatedMemory();
        mem += filename_cache_->getAllocatedMemory();

        return mem;
}

uint64_t SMBProtocol::getTotalAllocatedMemory() const {

        uint64_t mem = getAllocatedMemory();

	mem += compute_memory_used_by_maps();

        return mem;
}

uint64_t SMBProtocol::compute_memory_used_by_maps() const {

        uint64_t bytes = 0;

        std::for_each (filename_map_.begin(), filename_map_.end(), [&bytes] (PairStringCacheHits const &f) {
                bytes += f.first.size();
        });
        return bytes;
}

void SMBProtocol::increaseAllocatedMemory(int value) {

        info_cache_->create(value);
        filename_cache_->create(value);
}

void SMBProtocol::decreaseAllocatedMemory(int value) {

        info_cache_->destroy(value);
        filename_cache_->destroy(value);
}

void SMBProtocol::releaseCache() {

        if (FlowManagerPtr fm = flow_mng_.lock(); fm) {
                auto ft = fm->getFlowTable();

                std::ostringstream msg;
                msg << "Releasing " << name() << " cache";

                infoMessage(msg.str());

                uint64_t total_bytes_released = compute_memory_used_by_maps();
                uint64_t total_bytes_released_by_flows = 0;
                uint64_t reusable_bytes_by_flows = 0;
                uint32_t release_flows = 0;
		uint32_t release_filename = filename_map_.size();

                for (auto &flow: ft) {
                        if (SharedPointer<SMBInfo> info = flow->getSMBInfo(); info) {
                                total_bytes_released_by_flows += sizeof(info);

                                flow->layer7info.reset();
                                ++release_flows;
                                info_cache_->release(info);
                        }
                }
                // Some entries can be still on the maps and needs to be
                // retrieve to their existing caches
                for (auto &entry: filename_map_) {
			reusable_bytes_by_flows += entry.second.sc->size() * (entry.second.hits - 1);
                        releaseStringToCache(filename_cache_, entry.second.sc);
		}
		filename_map_.clear();

		data_time_ = boost::posix_time::microsec_clock::local_time();

                double cache_compression_rate = 0;

                if (reusable_bytes_by_flows > 0 )
                        cache_compression_rate = 100 - ((total_bytes_released * 100) / reusable_bytes_by_flows);

                msg.str("");
		msg << "Release " << release_filename << " file names, " << release_flows << " flows";
                msg << ", " << total_bytes_released + total_bytes_released_by_flows << " bytes";
                msg << ", " << reusable_bytes_by_flows << " reuse bytes";
                msg << ", compression rate " << cache_compression_rate << "%";
                infoMessage(msg.str());
        }
}

void SMBProtocol::releaseFlowInfo(Flow *flow) {

	if (auto info = flow->getSMBInfo(); info)
		info_cache_->release(info);
}

void SMBProtocol::attach_filename(SMBInfo *info, const boost::string_ref &name) {

	if (StringMap::iterator it = filename_map_.find(name); it != filename_map_.end()) {
		++(it->second).hits;
		info->filename = (it->second).sc;
	} else {
		if (SharedPointer<StringCache> name_ptr = filename_cache_->acquire(); name_ptr) {
			name_ptr->name(name.data(), name.length());
			info->filename = name_ptr;
			filename_map_.insert(std::make_pair(name_ptr->name(), name_ptr));
		}
        }
}

void SMBProtocol::create_request_directory_v1(SMBInfo *info, const uint8_t *payload, int length) {

        if (length >= (int)sizeof(smb_v1_create_file_request)) {
                const smb_v1_create_directory_request *hdr = reinterpret_cast<const smb_v1_create_directory_request*>(payload);
                int len = hdr->byte_count - 1;

                if (len <= (length - (int)sizeof(smb_v1_create_directory_request) + 1)) {
                        std::string filename;
                        const uint8_t *filename_payload = &hdr->data[0];

                        for (int i = 0; i < len; ++i) {
                                if (filename_payload[i] != 0x00)
                                        filename += (char)filename_payload[i];
                        }
                        if (filename.length() > 0) { // Something to attach
                                boost::string_ref file(filename.c_str(), filename.length());
                                attach_filename(info, file);
                        }
                }
        }
}

void SMBProtocol::create_request_file_v1(SMBInfo *info, const uint8_t *payload, int length) {

	if (length >= (int)sizeof(smb_v1_create_file_request)) {
		const smb_v1_create_file_request *hdr = reinterpret_cast<const smb_v1_create_file_request*>(payload);
		int len = hdr->byte_count - 1;

		if (len <= (length - (int)sizeof(smb_v1_create_file_request) + 1)) {
			std::string filename;
			const uint8_t *filename_payload = &hdr->data[0];

			for (int i = 0; i < len; ++i) {
				if (filename_payload[i] != 0x00) {
					filename += (char)filename_payload[i];
				}
			}
			if (filename.length() > 0) { // Something to attach
				boost::string_ref file(filename.c_str(), filename.length());
				attach_filename(info, file);
			}
		}
	}
}

void SMBProtocol::create_request_file_v2(SMBInfo *info, const uint8_t *payload, int length) {

	if (length >= (int)sizeof(smb_v2_create_file_request)) {
		const smb_v2_create_file_request *hdr = reinterpret_cast<const smb_v2_create_file_request*>(payload);
		int offset = hdr->filename_offset - (sizeof(smb_v2_header));
		if (offset < length) {
			std::string filename;
			const uint8_t *filename_payload = &payload[offset];
			for (int i = 0; i < hdr->filename_length; ++i) {
				if (filename_payload[i] != 0x00) {
					filename += (char)filename_payload[i];
				}
			}
			if (filename.length() > 0) { // Something to attach
				boost::string_ref file(filename.c_str(), filename.length());
				attach_filename(info, file);
			}
		}
	}
}

void SMBProtocol::open_and_request_v1(SMBInfo *info, const uint8_t *payload, int length) {

        if (length >= (int)sizeof(smb_v1_open_and_request)) {
        	const smb_v1_open_and_request *hdr = reinterpret_cast<const smb_v1_open_and_request*>(payload);

		if (hdr->byte_count + (int)sizeof(smb_v1_open_and_request) <= length) {
			if (hdr->byte_count > 0) {
				char *filename_payload = (char*)hdr->data;

				boost::string_ref file(filename_payload, hdr->byte_count);
				attach_filename(info, file);
			}
                }
        }
}

void SMBProtocol::update_command_type_v1(uint16_t cmd) {

	if ((cmd == SMB_CMD_OPEN_FILE)or(cmd == SMB_CMD_OPEN_ANDX))
		++total_open_files_;
	else if (cmd == SMB_CMD_CREATE_DIR)
		++total_create_dirs_;
	else if (cmd == SMB_CMD_DELETE_DIR)
		++total_delete_dirs_;
	else if (cmd == SMB_CMD_CREATE_FILE)
		++total_create_files_;
	else if (cmd == SMB_CMD_CLOSE_FILE)
		++total_close_files_;
	else if (cmd == SMB_CMD_FLUSH_FILES)
		++total_commit_files_;
	else if (cmd == SMB_CMD_DELETE_FILE)
		++total_delete_files_;
	else if (cmd == SMB_CMD_RENAME_FILE)
		++total_rename_files_;
	else if (cmd == SMB_CMD_GET_FILE_ATTR)
		++total_get_file_attribs_;
	else if (cmd == SMB_CMD_SET_FILE_ATTR)
		++total_set_file_attribs_;
        else if (cmd == SMB_CMD_TREE_DISC)
                ++total_tree_disconnects_;
        else if (cmd == SMB_CMD_NEGO_PROTO)
                ++total_negotiate_protocol_;
        else if (cmd == SMB_CMD_SESSION_SETUP)
                ++total_session_setups_;
        else if (cmd == SMB_CMD_LOGOFF)
                ++total_logoff_and_request_;
        else if (cmd == SMB_CMD_TREE_CONNECT)
                ++total_tree_connects_;
        else if ((cmd == SMB_CMD_TRANS)or(cmd == SMB_CMD_TRANS2))
                ++total_trans_;
        else if (cmd == SMB_CMD_READ)
                ++total_reads_;
        else if (cmd == SMB_CMD_NT_CREATE)
                ++total_nt_creates_;
        else if (cmd == SMB_CMD_WRITE_ANDX)
                ++total_writes_;
        else {
                ++total_others_;
	}
}

void SMBProtocol::update_command_type_v2(uint16_t cmd) {

	if (cmd == SMB2_CMD_CREATE_FILE)
		++total_create_files_;
	else if (cmd == SMB2_CMD_CLOSE_FILE)
		++total_close_files_;
	else if (cmd == SMB2_CMD_GET_INFO)
		++total_get_file_attribs_;
        else if (cmd == SMB2_CMD_TREE_CONNECT)
                ++total_tree_connects_;
        else if (cmd == SMB2_CMD_READ_FILE)
                ++total_reads_;
        else if (cmd == SMB2_CMD_WRITE_FILE)
                ++total_writes_;
        else {
                ++total_others_;
	}
}

void SMBProtocol::processFlow(Flow *flow) {

	CPUCycle cycles(&total_cpu_cycles_);
	int length = flow->packet->getLength();
	total_bytes_ += length;
	++total_packets_;
	++flow->total_packets_l7;

	current_flow_ = flow;

	if (length >= header_size) {
       		SharedPointer<SMBInfo> info = flow->getSMBInfo();
       		if (!info) {
               		if (info = info_cache_->acquire(); !info) {
				logFailCache(info_cache_->name(), flow);
				return;
       			}
               		flow->layer7info = info;
       		}

		const uint8_t *payload = flow->packet->getPayload();

		if (is_minimal_smb_header(payload)) {
			uint8_t smb_version = payload[4];
			uint16_t cmd = 0;

			if (smb_version == 0xFF) { // Version 1
				const smb_v1_header *hdr = reinterpret_cast<const smb_v1_header*>(&payload[4]);
				cmd = hdr->cmd;
				int offset = 4 + sizeof(smb_v1_header);
				update_command_type_v1(hdr->cmd);
				if (cmd == SMB_CMD_OPEN_ANDX)
					open_and_request_v1(info.get(), &payload[offset], length - offset);
				else if (cmd == SMB_CMD_CREATE_FILE)
					create_request_file_v1(info.get(), &payload[offset], length - offset);
				else if (cmd == SMB_CMD_CREATE_DIR)
					create_request_directory_v1(info.get(), &payload[offset], length - offset);

			} else if (smb_version == 0xFE) { // Version 2
				const smb_v2_header *hdr = reinterpret_cast<const smb_v2_header*>(&payload[4]);
				cmd = hdr->cmd;
				update_command_type_v2(cmd);
				int offset = 4 + sizeof(smb_v2_header);
				if (cmd == SMB2_CMD_CREATE_FILE)
					create_request_file_v2(info.get(), &payload[offset], length - offset);

			} else {
				++total_events_;
                		// Malformed header packet
                		if (current_flow_->getPacketAnomaly() == PacketAnomalyType::NONE)
                        		current_flow_->setPacketAnomaly(PacketAnomalyType::SMB_BOGUS_HEADER);

                		anomaly_->incAnomaly(PacketAnomalyType::SMB_BOGUS_HEADER);
                		return;
			}

			info->setCommand(cmd);
		} else {
			// TODO Current transfer of files , may be regex?
		}
	} else {
		++total_events_;
                // Malformed header packet
                if (current_flow_->getPacketAnomaly() == PacketAnomalyType::NONE)
                	current_flow_->setPacketAnomaly(PacketAnomalyType::SMB_BOGUS_HEADER);
	}
}

uint32_t SMBProtocol::getTotalCacheMisses() const {

        uint32_t miss = 0;

        miss = info_cache_->getTotalFails();
        miss += filename_cache_->getTotalFails();

        return miss;
}

void SMBProtocol::statistics(std::basic_ostream<char>& out, int level, int32_t limit) const {

	showStatisticsHeader(out, level);

	if (level > 3) {
		out << "\t" << "Total create dirs:      " << std::setw(10) << total_create_dirs_ << "\n"
			<< "\t" << "Total delete dirs:      " << std::setw(10) << total_delete_dirs_ << "\n"
			<< "\t" << "Total open files:       " << std::setw(10) << total_open_files_ << "\n"
			<< "\t" << "Total create files:     " << std::setw(10) << total_create_files_ << "\n"
			<< "\t" << "Total close files:      " << std::setw(10) << total_close_files_ << "\n"
			<< "\t" << "Total commit files:     " << std::setw(10) << total_commit_files_ << "\n"
			<< "\t" << "Total delete files:     " << std::setw(10) << total_delete_files_ << "\n"
			<< "\t" << "Total rename files:     " << std::setw(10) << total_rename_files_ << "\n"
			<< "\t" << "Total get attr files:   " << std::setw(10) << total_get_file_attribs_ << "\n"
			<< "\t" << "Total set attr files:   " << std::setw(10) << total_set_file_attribs_ << "\n"
			<< "\t" << "Total tree disconnects: " << std::setw(10) << total_tree_disconnects_ << "\n"
			<< "\t" << "Total negotiate proto:  " << std::setw(10) << total_negotiate_protocol_ << "\n"
			<< "\t" << "Total session setups:   " << std::setw(10) << total_session_setups_ << "\n"
			<< "\t" << "Total logoffs:          " << std::setw(10) << total_logoff_and_request_ << "\n"
			<< "\t" << "Total tree connects:    " << std::setw(10) << total_tree_connects_ << "\n"
			<< "\t" << "Total trans:            " << std::setw(10) << total_trans_ << "\n"
			<< "\t" << "Total reads:            " << std::setw(10) << total_reads_ << "\n"
			<< "\t" << "Total writes:           " << std::setw(10) << total_writes_ << "\n"
			<< "\t" << "Total nt creates:       " << std::setw(10) << total_nt_creates_ << "\n"
			<< "\t" << "Total others:           " << std::setw(10) << total_others_ << std::endl;
	}
	if ((level > 5)and(flow_forwarder_.lock()))
		flow_forwarder_.lock()->statistics(out);
	if (level > 3) {
		info_cache_->statistics(out);
		filename_cache_->statistics(out);
		if (level > 4)
			filename_map_.show(out, "\t", limit);
	}
}

void SMBProtocol::statistics(Json &out, int level) const {

	showStatisticsHeader(out, level);

        if (level > 3) {
		Json j;

		j["create_dirs"] = total_create_dirs_;
		j["delete_dirs"] = total_delete_dirs_;
		j["open_files"] = total_open_files_;
		j["create_files"] = total_create_files_;
		j["close_files"] = total_close_files_;
		j["commit_files"] = total_commit_files_;
		j["delete_files"] = total_delete_files_;
		j["rename_files"] = total_rename_files_;
		j["get_attr_files"] = total_get_file_attribs_;
		j["set_attr_files"] = total_set_file_attribs_;
		j["tree_disconnects"] = total_tree_disconnects_;
		j["negotiate_protos"] = total_negotiate_protocol_;
		j["session_setups"] = total_session_setups_;
		j["logoffs"] = total_logoff_and_request_;
		j["tree_connects"] = total_tree_connects_;
		j["trans"] = total_trans_;
		j["reads"] = total_reads_;
		j["writes"] = total_writes_;
		j["nt_creates"] = total_nt_creates_;
		j["others"] = total_others_;

		out["methods"] = j;
        }
}

CounterMap SMBProtocol::getCounters() const {
       	CounterMap cm;

        cm.addKeyValue("packets", total_packets_);
        cm.addKeyValue("bytes", total_bytes_);

	// Specific SMB counters
	cm.addKeyValue("create dirs", total_create_dirs_);
	cm.addKeyValue("delete dirs", total_delete_dirs_);
	cm.addKeyValue("open files", total_open_files_);
	cm.addKeyValue("create files", total_create_files_);
	cm.addKeyValue("close files", total_close_files_);
	cm.addKeyValue("commit files", total_commit_files_);
	cm.addKeyValue("delete files", total_delete_files_);
	cm.addKeyValue("rename files", total_rename_files_);
	cm.addKeyValue("get attr files", total_get_file_attribs_);
	cm.addKeyValue("set attr files", total_set_file_attribs_);
	cm.addKeyValue("tree disconnects", total_tree_disconnects_);
	cm.addKeyValue("negotiate proto", total_negotiate_protocol_);
	cm.addKeyValue("session setups", total_session_setups_);
	cm.addKeyValue("logoffs", total_logoff_and_request_);
	cm.addKeyValue("tree connects", total_tree_connects_);
	cm.addKeyValue("trans", total_trans_);
	cm.addKeyValue("reads", total_reads_);
	cm.addKeyValue("writes", total_writes_);
	cm.addKeyValue("nt creates", total_nt_creates_);
	cm.addKeyValue("others", total_others_);

        return cm;
}

#if defined(PYTHON_BINDING) || defined(RUBY_BINDING)
#if defined(PYTHON_BINDING)
boost::python::dict SMBProtocol::getCacheData(const std::string &name) const {
#elif defined(RUBY_BINDING)
VALUE SMBProtocol::getCacheData(const std::string &name) const {
#endif
        if (boost::iequals(name, "filename"))
		return addMapToHash(filename_map_);

	StringMap empty {"", ""};

        return addMapToHash(empty);
}

#if defined(PYTHON_BINDING)
SharedPointer<Cache<StringCache>> SMBProtocol::getCache(const std::string &name) {

        if (boost::iequals(name, "filename"))
                return filename_cache_;

        return nullptr;
}

#endif

#endif

void SMBProtocol::statistics(Json &out, const std::string &map_name, int32_t limit) const {

        if (boost::iequals(map_name, "filenames")) {
                filename_map_.show(out, limit);
        }
}

void SMBProtocol::resetCounters() {

	reset();

        total_create_dirs_ = 0;
        total_delete_dirs_ = 0;
        total_open_files_ = 0;
        total_create_files_ = 0;
        total_close_files_ = 0;
        total_commit_files_ = 0;
        total_delete_files_ = 0;
        total_rename_files_ = 0;
        total_get_file_attribs_ = 0;
        total_set_file_attribs_ = 0;
        total_tree_disconnects_ = 0;
        total_negotiate_protocol_ = 0;
        total_session_setups_ = 0;
        total_logoff_and_request_ = 0;
        total_tree_connects_ = 0;
        total_trans_ = 0;
        total_reads_ = 0;
        total_writes_ = 0;
        total_nt_creates_ = 0;
        total_others_ = 0;
        total_events_ = 0;
}

} // namespace aiengine
